﻿// dllmain.cpp : 定义 DLL 应用程序的入口点。
#include "pch.h"
#include <msi.h>
#include <msiquery.h>
#include <WinUser.h>

BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
                     )
{
    switch (ul_reason_for_call)
    {
    case DLL_PROCESS_ATTACH:
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}

LPWSTR copyEnvironmentVariable(LPCWSTR variable)
{
	DWORD length = GetEnvironmentVariable(variable, 0, 0);
	if (!length)
		return nullptr;
	LPWSTR buffer = (LPWSTR)HeapAlloc(GetProcessHeap(), 0, (length + 10) * sizeof(wchar_t));
	if (!GetEnvironmentVariable(variable, buffer, length) || !buffer[0])
		return nullptr;
	return buffer;
}

DWORD AdjustProcessToken(HANDLE hProcess, LPCTSTR lpPrivilegeName)
{
	HANDLE				hToken;		//令牌句柄
	TOKEN_PRIVILEGES	tp;
	tp.PrivilegeCount = 1;
	tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	if (!OpenProcessToken(hProcess, TOKEN_ADJUST_PRIVILEGES, &hToken)) {
		return GetLastError();
	}
	if (!LookupPrivilegeValue(NULL, lpPrivilegeName, &tp.Privileges[0].Luid)) {
		return GetLastError();
	}
	if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), NULL, NULL)) {
		return GetLastError();
	}
	return 0;
}

UINT __stdcall AfterCommit(MSIHANDLE hModule)
{
	bool bRet = false; HKEY hSubKey = nullptr; do {
		wchar_t target[MAX_PATH]; DWORD len = MAX_PATH;
		if (MsiGetProperty(hModule, L"CustomActionData", target, &len) != ERROR_SUCCESS) break;
		wcscat_s(target, L"Apple Application Support");
		AdjustProcessToken(GetCurrentProcess(), SE_DEBUG_NAME);
		if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, L"System\\CurrentControlSet\\Control\\Session Manager\\Environment", 0, KEY_ALL_ACCESS, &hSubKey) != ERROR_SUCCESS) break;		// 打开注册表
		wchar_t regPath[32768]; DWORD typeSZ = REG_SZ; DWORD cbName = sizeof(regPath);
		if (RegQueryValueEx(hSubKey, L"Path", 0, &typeSZ, (LPBYTE)regPath, &cbName) != ERROR_SUCCESS) break;
		if (!regPath[0])
			wcscpy_s(regPath, target);
		else if (!wcsstr(regPath, target))
		{
			wcscat_s(regPath, L";"); wcscat_s(regPath, target);
		}
		else
		{
			bRet = true; break;
		}
		if (RegSetValueEx(hSubKey, L"Path", 0, REG_SZ, (LPBYTE)regPath, (DWORD)(wcslen(regPath) + 1) * sizeof(wchar_t)) != ERROR_SUCCESS) break;
		DWORD_PTR dwReturnValue;
		SendMessageTimeout(HWND_BROADCAST, WM_SETTINGCHANGE, 0, (LPARAM)L"Environment", SMTO_ABORTIFHUNG, 5000, &dwReturnValue);
		bRet = true;
	} while (0);
	if (!hSubKey) RegCloseKey(hSubKey);
	if (!bRet)
		MessageBox(nullptr, L"请将安装目录下的Apple Application Support子目录加到系统的PATH环境变量",
			L"设置环境变量失败", MB_ICONEXCLAMATION);
	return ERROR_SUCCESS;
}
